import numpy as np
import torch
from torch import nn
import torch.nn.functional as F
from torch.autograd import Variable
from torchvision.datasets import CIFAR10
from utils import train
from torch.utils.data import DataLoader

def conv3x3(in_channel,out_channel,stride=1):
    return nn.Conv2d(in_channel, out_channel, 3, stride=stride, padding=1, bias=False)

class residual_block(nn.Module):
    def __init__(self,in_channel,out_channel,same_shape=True):
        super(residual_block,self).__init__()
        self.same_shape = same_shape
        stride = 1 if self.same_shape else 2

        self.conv1 = conv3x3(in_channel, out_channel, stride = stride)
        self.bn1 = nn.BatchNorm2d(out_channel)
        self.conv2 = conv3x3(out_channel, out_channel)
        self.bn2 = nn.BatchNorm2d(out_channel)
        if not self.same_shape:
            self.conv3 = nn.Conv2d(in_channel, out_channel, 1, stride= stride)
        
    def forward(self,x):
        out = self.conv1(x)
        out = F.relu(self.bn1(out),True)
        out = self.conv2(out)
        out = F.relu(self.bn2(out),True)

        if not self.same_shape:
            x = self.conv3(x)
        return F.relu(x+out, True)
#test_shape    
# test_net = residual_block(3,32,False)
# test_x = Variable(torch.zeros(1,3,96,96))
# print('input:{}'.format(test_x.shape))
# test_y = test_net(test_x)
# print('output:{}'.format(test_y.shape))

class res_net(nn.Module):
    def __init__(self, in_channel, num_classes, verbose = False):
        super(res_net,self).__init__()
        self.verbose = verbose

        self.block1 = nn.Conv2d(in_channel, 64, 7, 2)

        self.block2 = nn.Sequential(
            nn.MaxPool2d(3,2),
            residual_block(64,64),
            residual_block(64,64)
        )

        self.block3 = nn.Sequential(
            residual_block(64,128,False),
            residual_block(128,128)
        )

        self.block4 = nn.Sequential(
            residual_block(128,256,False),
            residual_block(256,256)
        )

        self.block5 = nn.Sequential(
            residual_block(256,512,False),
            residual_block(512,512),
            nn.AvgPool2d(3)
        )

        self.classifier = nn.Linear(512, num_classes)

    def forward(self,x):
        x = self.block1(x)
        if self.verbose:
            print('block 1 output {}'.format(x.shape))
        x = self.block2(x)
        if self.verbose:
            print('block 2 output {}'.format(x.shape))
        x = self.block3(x)
        if self.verbose:
            print('block 3 output {}'.format(x.shape))
        x = self.block4(x)
        if self.verbose:
            print('block 4 output {}'.format(x.shape))       
        x = self.block5(x)
        if self.verbose:
            print('block 5 output {}'.format(x.shape))   
        x = x.view(x.shape[0],-1)
        x = self.classifier(x)
        return x
    
# test_net = res_net(3, 10, True)
# test_x = Variable(torch.zeros(1, 3, 96, 96))
# test_y = test_net(test_x)
# print('output: {}'.format(test_y.shape))

def data_tf(x):
    x = x.resize((96,96),2)
    x = np.array(x,dtype='float32')/255
    x = (x-0.5)/0.5
    x = x.transpose((2,0,1))
    x = torch.from_numpy(x)
    return x


train_set = CIFAR10('./data', train=True, transform=data_tf, download=True)
train_data = DataLoader(train_set, batch_size=64, shuffle=True)
# test_set = CIFAR10('./data', train=False, transform=data_tf, download=True)
# test_train = DataLoader(test_set, batch_size=64, shuffle=False)
# a,label = next(iter(train_data))
# print(a.shape)
# print(label.shape)
net = res_net(3,10)
optimizer = torch.optim.SGD(net.parameters(),lr=0.01)
cre = nn.CrossEntropyLoss()


